package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"
	"time"
)

type ctxKey int

const (
	backupTargetsKey ctxKey = iota
)

type Proxy struct {
	state         *State
	rp            *httputil.ReverseProxy
	transport     *http.Transport
	balancers     map[string]LoadBalancer // domain -> load balancer
	balancersMu   sync.RWMutex
	baseTLSConfig *tls.Config
	targetCache   map[string][]*url.URL
	targetCacheMu sync.RWMutex
}

func NewProxy(state *State) *Proxy {
	// A highly reused Transport is key for performance and low memory
	// 修复 Go 1.22 TLS 握手失败问题：显式配置包含所有密码套件的 TLS 配置
	// Go 1.22 从默认列表中移除了基于 RSA 密钥交换的密码套件，需要手动指定
	// 参考: https://github.com/golang/go/issues/66512
	//
	// 注意：使用动态 TLS 配置，根据目标地址调整（参考 nginx 的实现）
	// - 对于 IP 地址，不发送 SNI（RFC 6066 不允许对 IP 地址使用 SNI）
	// - 对于本地服务器（127.0.0.1, localhost），跳过证书验证
	baseTLSConfig := &tls.Config{
		CipherSuites: []uint16{
			// TLS 1.0 - 1.2 cipher suites (包括基于 RSA 密钥交换的密码套件)
			tls.TLS_RSA_WITH_RC4_128_SHA,
			tls.TLS_RSA_WITH_3DES_EDE_CBC_SHA,
			tls.TLS_RSA_WITH_AES_128_CBC_SHA,
			tls.TLS_RSA_WITH_AES_256_CBC_SHA,
			tls.TLS_RSA_WITH_AES_128_CBC_SHA256,
			tls.TLS_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_RC4_128_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_RC4_128_SHA,
			tls.TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,
			// TLS 1.3 cipher suites (不可配置，但列出以保持完整性)
			// TLS 1.3 会自动使用这些密码套件
		},
		MinVersion: tls.VersionTLS12,
	}

	// 使用 DialTLSContext 来动态配置 TLS（根据目标地址）
	dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 60 * time.Second}

	tr := &http.Transport{
		Proxy:       http.ProxyFromEnvironment,
		DialContext: dialer.DialContext,
		DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			// 解析目标地址
			host, _, err := net.SplitHostPort(addr)
			if err != nil {
				host = addr
			}

			// 判断是否是 IP 地址
			isIP := net.ParseIP(host) != nil
			isLocalhost := host == "localhost" || host == "127.0.0.1" || host == "::1"

			// 创建针对此连接的 TLS 配置
			tlsConfig := baseTLSConfig.Clone()
			if isLocalhost || isIP {
				// 对于本地服务器或 IP 地址，跳过证书验证（参考 nginx 的 proxy_ssl_verify off）
				tlsConfig.InsecureSkipVerify = true
			}
			if isIP {
				// 对于 IP 地址，不发送 SNI（RFC 6066 不允许，参考 nginx 实现）
				tlsConfig.ServerName = ""
			} else {
				// 对于域名，使用主机名作为 SNI
				tlsConfig.ServerName = host
			}

			// 建立 TCP 连接
			conn, err := dialer.DialContext(ctx, network, addr)
			if err != nil {
				return nil, err
			}

			// 建立 TLS 连接
			tlsConn := tls.Client(conn, tlsConfig)
			if err := tlsConn.HandshakeContext(ctx); err != nil {
				conn.Close()
				return nil, err
			}

			return tlsConn, nil
		},
		ForceAttemptHTTP2:     true,
		MaxIdleConns:          1024,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
		DisableCompression:    false,
		// 设置基础 TLS 配置（用于非 DialTLSContext 的情况）
		TLSClientConfig: baseTLSConfig,
	}

	p := &Proxy{
		state:         state,
		transport:     tr,
		balancers:     make(map[string]LoadBalancer),
		baseTLSConfig: baseTLSConfig,
		targetCache:   make(map[string][]*url.URL),
	}

	director := func(req *http.Request) {
		host := stripPort(req.Host)
		up, ok := p.state.GetUpstream(host)
		if !ok {
			log.Printf("proxy: no upstream found for domain: %s", host)
			return
		}

		targets := p.buildTargetChain(host, up)
		if len(targets) == 0 {
			log.Printf("proxy: no valid target found for %s (servers: %d, target: %s)", host, len(up.Servers), up.Target)
			return
		}

		tgt := targets[0]
		if len(targets) > 1 {
			p.attachBackupTargets(req, targets[1:])
		}

		// rewrite request
		req.URL.Scheme = tgt.Scheme
		req.URL.Host = tgt.Host
		// optional: preserve original Host header or set to upstream host
		req.Host = tgt.Host

		log.Printf("proxy: forwarding %s %s to %s://%s%s (backups: %d)", req.Method, req.URL.Path, tgt.Scheme, tgt.Host, req.URL.Path, len(targets)-1)
		// WebSocket and HTTP/1.1 are handled by ReverseProxy transparently
	}

	errorLogger := log.New(log.Writer(), "proxy ", log.LstdFlags)

	p.rp = &httputil.ReverseProxy{
		Director:  director,
		Transport: tr,
		ErrorLog:  errorLogger,
		ModifyResponse: func(resp *http.Response) error {
			return nil
		},
		ErrorHandler: func(rw http.ResponseWriter, r *http.Request, err error) {
			log.Printf("proxy: upstream error for %s %s: %v", r.Method, r.URL.Path, err)

			backups := getBackupTargets(r)
			if len(backups) > 0 {
				for i, nextTarget := range backups {
					log.Printf("proxy: retrying with backup server %d/%d: %s://%s", i+1, len(backups), nextTarget.Scheme, nextTarget.Host)

					remaining := backups[i+1:]
					newCtx := context.WithValue(r.Context(), backupTargetsKey, remaining)
					newReq := r.Clone(newCtx)
					newReq.URL.Scheme = nextTarget.Scheme
					newReq.URL.Host = nextTarget.Host
					newReq.Host = nextTarget.Host

					resp, retryErr := p.transport.RoundTrip(newReq)
					if retryErr == nil && resp != nil {
						for k, v := range resp.Header {
							for _, val := range v {
								rw.Header().Add(k, val)
							}
						}
						rw.WriteHeader(resp.StatusCode)
						if resp.Body != nil {
							defer resp.Body.Close()
							io.Copy(rw, resp.Body)
						}
						log.Printf("proxy: successfully retried with backup server %s://%s", nextTarget.Scheme, nextTarget.Host)
						return
					}

					log.Printf("proxy: backup server %s://%s also failed: %v", nextTarget.Scheme, nextTarget.Host, retryErr)
				}
			}

			// 所有服务器都失败了
			http.Error(rw, fmt.Sprintf("upstream error: all servers failed, last error: %v", err), http.StatusBadGateway)
		},
	}
	return p
}

func (p *Proxy) attachBackupTargets(req *http.Request, backups []*url.URL) {
	if len(backups) == 0 {
		return
	}
	ctx := context.WithValue(req.Context(), backupTargetsKey, backups)
	*req = *req.WithContext(ctx)
}

func getBackupTargets(r *http.Request) []*url.URL {
	if v := r.Context().Value(backupTargetsKey); v != nil {
		if targets, ok := v.([]*url.URL); ok {
			return targets
		}
	}
	return nil
}

func (p *Proxy) buildTargetChain(domain string, up Upstream) []*url.URL {
	servers := p.cachedServerURLs(domain, up)
	if len(servers) == 0 {
		return nil
	}
	chain := make([]*url.URL, 0, len(servers)+1)
	seen := make(map[string]struct{}, len(servers)+1)
	if lb := p.getOrCreateBalancer(domain, up); lb != nil {
		if primary, _ := lb.Select(); primary != nil {
			key := primary.String()
			chain = append(chain, primary)
			seen[key] = struct{}{}
		}
	}
	for _, srv := range servers {
		key := srv.String()
		if _, exists := seen[key]; exists {
			continue
		}
		seen[key] = struct{}{}
		chain = append(chain, srv)
	}
	return chain
}

func (p *Proxy) cachedServerURLs(domain string, up Upstream) []*url.URL {
	p.targetCacheMu.RLock()
	urls, ok := p.targetCache[domain]
	p.targetCacheMu.RUnlock()
	if ok && len(urls) > 0 {
		return urls
	}
	urls = make([]*url.URL, 0, len(up.Servers))
	for _, srv := range up.Servers {
		if u, err := url.Parse(srv.Target); err == nil {
			urls = append(urls, u)
		}
	}
	if len(urls) == 0 && up.Target != "" {
		if u, err := url.Parse(up.Target); err == nil {
			urls = append(urls, u)
		}
	}
	if len(urls) > 0 {
		p.targetCacheMu.Lock()
		p.targetCache[domain] = urls
		p.targetCacheMu.Unlock()
	}
	return urls
}

// getOrCreateBalancer 获取或创建负载均衡器
func (p *Proxy) getOrCreateBalancer(domain string, up Upstream) LoadBalancer {
	p.balancersMu.RLock()
	lb, exists := p.balancers[domain]
	p.balancersMu.RUnlock()

	if exists {
		return lb
	}

	// 创建新的负载均衡器
	lb, err := CreateLoadBalancer(up)
	if err != nil || lb == nil {
		return nil
	}

	p.balancersMu.Lock()
	p.balancers[domain] = lb
	p.balancersMu.Unlock()

	return lb
}

// invalidateBalancer 使负载均衡器失效（当配置更新时调用）
func (p *Proxy) invalidateBalancer(domain string) {
	p.balancersMu.Lock()
	delete(p.balancers, domain)
	p.balancersMu.Unlock()
	p.targetCacheMu.Lock()
	delete(p.targetCache, domain)
	p.targetCacheMu.Unlock()
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if _, ok := p.state.GetUpstream(r.Host); !ok {
		http.NotFound(w, r)
		return
	}
	p.rp.ServeHTTP(w, r)
}
